package partition;

import network.Device;
import network.Network;
import network.Port;
import network.Rule;
import verifier.JiffyVerifier;
import verifier.Ports.PersistentPorts;
import verifier.util.Changes;

import java.util.*;

public class Partition {
    public static int cnt;
    public int uid;
    protected String name;
    protected Network network;
    private List<LocalTask> localTasks;
    protected JiffyVerifier verifier;
    private HashMap<Device, ArrayList<LocalTask>> ingressToLocalTasks;
    private HashMap<Port, ArrayList<LocalTask>> egressToLocalTasks;

    public Partition(String name, Network network) {
        this.name = name;
        this.network = network;
        this.localTasks = new ArrayList<>();
//        this.apKeep = new APKeep(network);
//        this.verifier = new Verifier(network);
        this.ingressToLocalTasks = new HashMap<>();
        this.egressToLocalTasks = new HashMap<>();
        for (Device device : network.getAllDevices()) {
            device.setPartition(this);
        }
    }

    public void addLocalTask(LocalTask localTask) {
        this.localTasks.add(localTask);
        this.ingressToLocalTasks.computeIfAbsent(localTask.getIngress(), k -> new ArrayList<>());
        this.ingressToLocalTasks.get(localTask.getIngress()).add(localTask);
        this.egressToLocalTasks.computeIfAbsent(localTask.getEgress(), k -> new ArrayList<>());
        this.egressToLocalTasks.get(localTask.getEgress()).add(localTask);
    }

    public ArrayList<LocalTask> getLocalTasksByIngress(Device ingress) {
        return this.ingressToLocalTasks.get(ingress);
    }

    public ArrayList<LocalTask> getLocalTasksByEgress(Port egress) {
        return this.egressToLocalTasks.get(egress);
    }

//    public void run() {
//        long s = System.nanoTime();
//        for (Device device : this.getNetwork().getAllDevices()) {
//            for (Rule rule : device.getInitialRules()) {
//                this.addRule(rule);
//            }
//        }
//        System.out.println("model update of " + this.name + " takes: " + (System.nanoTime() - s));
//    }

    public void runLocalTasks() {
        for (LocalTask localTask : this.localTasks) {
//            long start = System.nanoTime();
            localTask.run(this.verifier);
//            System.out.println("localtask " + localTask.getIngress().getName() + "->" + localTask.getEgress().getName() + " takes: " + (System.nanoTime() - start));
        }
    }

    public void addRule(Rule rule) {
        Changes changes = this.verifier.insertMiniBatch(new ArrayList<>(Arrays.asList(rule)));
        changes.aggrBDDs();
        this.verifier.update(changes);
    }

    /**
     * Evaluate mini-batch
     *
     * @param M #updates per-mini-batch
     * @return Model update time (us per-update)
     */
    public double test(int M) {
        /**
         Important: each APKeep instance re-labels every device when initialized.
         Given the Device instances are shared, only one APKeep instance can work at the same time.
         */
        this.verifier = new JiffyVerifier(network, new PersistentPorts());

        int N = 0;
        ArrayList<ArrayList<Rule>> ruless = new ArrayList<>();

        int cnt = 0, r = 0;
        for (Device device : this.getNetwork().getAllDevices()) {
            for (Rule rule : device.getInitialRules()) {
                if (cnt % M == 0) {
                    ruless.add(new ArrayList<>());
                    N++;
                }
                ruless.get(N - 1).add(rule);
                cnt++;
            }
        }

        // sort batches
        // Comparator<ArrayList<Rule>> comp = (ArrayList<Rule> lhs, ArrayList<Rule> rhs) -> rhs.get(0).getPriority() - lhs.get(0).getPriority();
        // ruless.sort(comp);

        System.out.println("# mini-batches " + ruless.size());
        long s = System.nanoTime();
        double tot = 0;
        // ArrayList<Double> times = new ArrayList<>();

        boolean flag1 = false, flag2 = false;
        for (ArrayList<Rule> rules : ruless) {
            if (flag1) System.out.println();
            if (flag1) System.out.println("===========================================");

            if (flag1) System.out.println("#rules in this mini-batch " + rules.size());
            double t0 = System.nanoTime();
            Changes changes = verifier.insertMiniBatch(rules);
            double t1 = System.nanoTime();
            if (flag1) System.out.println("Rule => change " + ((t1 - t0) / rules.size() / 1000) + " us per update");

//            if (flag1) System.out.println("Remaining BDDs: " + (BDDEngine.refCnt - BDDEngine.defCnt));

            if (flag1) System.out.println("# changes (raw) " + changes.aggr0Size());
            if (flag1) System.out.println("# changes (step-1-aggr) " + changes.aggr1Size());
            changes.aggrBDDs();
            if (flag1) System.out.println("# changes (step-2-aggr) " + changes.aggr2Size());

            double t2 = System.nanoTime();
            verifier.update(changes);
            double t3 = System.nanoTime();
            if (flag1) System.out.println("change => ECs " + ((t3 - t2) / rules.size() / 1000) + " us per update");
            // times.add((t3 - t2) / rules.size() / 1000);
            tot += t3 - t2;


            if (flag1) System.out.println("Total #EC: " + verifier.predSize());
        }

        if (flag2) {
            System.out.println("finish " + this.name + ": " + ((System.nanoTime() - s) / cnt / 1000) + " us per-update");
            System.out.println();
//            System.out.println("Total BDD : " + (BDDEngine.tot / cnt / 1000) + " us per-update");
//            System.out.println("BDD time/operation ratio : " + (BDDEngine.tot / BDDEngine.refCnt) + " ns");
//            System.out.println("Remaining BDDs: " + (BDDEngine.refCnt - BDDEngine.defCnt));
            System.out.println("Total #EC: " + verifier.predSize());
            // PrintWriter writer = new PrintWriter("times.log", "UTF-8");
            // for (Double t : times) writer.println(t);
            // writer.close();
            System.out.println("Change => EC: " + (tot / cnt / 1000) + " us per-update");
            System.exit(0);
        }

//        verifier = null;
//        System.gc();
        return (tot / cnt / 1000);
    }

    public void run(List<Rule> rules, boolean useFMT) {
        this.verifier = new JiffyVerifier(network, new PersistentPorts());
        if (useFMT) {
            Changes changes = verifier.insertMiniBatch(rules);
            changes.aggrBDDs();
            this.verifier.update(changes);
        } else {
            for (Rule rule : rules) {
                Changes changes = verifier.insertMiniBatch(new ArrayList<>(Collections.singletonList(rule)));
                changes.aggrBDDs();
                this.verifier.update(changes);
            }
        }
    }


    /**
     * Evaluate mini-batch
     *
     * @param M #updates per-mini-batch
     * @return Model update time (us per-update)
     */
    public double runSubnet(int M, Collection<Rule> initRules) {
        /**
         Important: each APKeep instance re-labels every device when initialized.
         Given the Device instances are shared, only one APKeep instance can work at the same time.
         */
        this.verifier = new JiffyVerifier(network, new PersistentPorts());

        int N = 0;
        ArrayList<ArrayList<Rule>> ruless = new ArrayList<>();

        int cnt = 0, r = 0;
        for (Rule rule : initRules) {
            if (cnt % M == 0) {
                ruless.add(new ArrayList<>());
                N++;
            }
            ruless.get(N - 1).add(rule);
            cnt++;
        }

        // sort batches
        // Comparator<ArrayList<Rule>> comp = (ArrayList<Rule> lhs, ArrayList<Rule> rhs) -> rhs.get(0).getPriority() - lhs.get(0).getPriority();
        // ruless.sort(comp);

//        System.out.println("# mini-batches " + ruless.size());
        long s = System.nanoTime();
        double tot = 0;
        // ArrayList<Double> times = new ArrayList<>();

        boolean flag1 = false, flag2 = false;
        for (ArrayList<Rule> rules : ruless) {
            if (flag1) System.out.println();
            if (flag1) System.out.println("===========================================");

            if (flag1) System.out.println("#rules in this mini-batch " + rules.size());
            double t0 = System.nanoTime();
            Changes changes = verifier.insertMiniBatch(rules);
            double t1 = System.nanoTime();
            if (flag1) System.out.println("Rule => change " + ((t1 - t0) / rules.size() / 1000) + " us per update");

//            if (flag1) System.out.println("Remaining BDDs: " + (BDDEngine.refCnt - BDDEngine.defCnt));

            if (flag1) System.out.println("# changes (raw) " + changes.aggr0Size());
            if (flag1) System.out.println("# changes (step-1-aggr) " + changes.aggr1Size());
            changes.aggrBDDs();
            if (flag1) System.out.println("# changes (step-2-aggr) " + changes.aggr2Size());

            double t2 = System.nanoTime();
            verifier.update(changes);
            double t3 = System.nanoTime();
            if (flag1) System.out.println("change => ECs " + ((t3 - t2) / rules.size() / 1000) + " us per update");
            // times.add((t3 - t2) / rules.size() / 1000);
            tot += t3 - t2;


            if (flag1) System.out.println("Total #EC: " + verifier.predSize());
        }
        System.out.println("Total #EC: " + verifier.predSize());
        return (tot / cnt / 1000);
    }

    public double runSubnetRuleBatch(ArrayList<ArrayList<Rule>> ruless) {
        /**
         Important: each APKeep instance re-labels every device when initialized.
         Given the Device instances are shared, only one APKeep instance can work at the same time.
         */
        this.verifier = new JiffyVerifier(network, new PersistentPorts());

        boolean flag1 = false, flag2 = false;
        for (ArrayList<Rule> rules : ruless) {
            if (flag1) System.out.println();
            if (flag1) System.out.println("===========================================");

            if (flag1) System.out.println("#rules in this mini-batch " + rules.size());
            double t0 = System.nanoTime();
            Changes changes = verifier.insertMiniBatch(rules);
            double t1 = System.nanoTime();
            if (flag1) System.out.println("Rule => change " + ((t1 - t0) / rules.size() / 1000) + " us per update");

//            if (flag1) System.out.println("Remaining BDDs: " + (BDDEngine.refCnt - BDDEngine.defCnt));

            if (flag1) System.out.println("# changes (raw) " + changes.aggr0Size());
            if (flag1) System.out.println("# changes (step-1-aggr) " + changes.aggr1Size());
            changes.aggrBDDs();
            if (flag1) System.out.println("# changes (step-2-aggr) " + changes.aggr2Size());

            double t2 = System.nanoTime();
            verifier.update(changes);
            double t3 = System.nanoTime();
            if (flag1) System.out.println("change => ECs " + ((t3 - t2) / rules.size() / 1000) + " us per update");
            // times.add((t3 - t2) / rules.size() / 1000);

            if (flag1) System.out.println("Total #EC: " + verifier.predSize());
        }
//        System.out.println("Total #EC: " + verifier.predSize());
        return 0;
    }

    public String getName() {
        return name;
    }

    public Network getNetwork() {
        return network;
    }

    public List<LocalTask> getLocalTasks() {
        return localTasks;
    }

    public JiffyVerifier getVerifier() {
        return this.verifier;
    }
}
